import os
import random
import plotly.graph_objects as go

from typing import Any, List, Dict
from data_process.structure import Point3D


class PointCloudVisualizer:
    @staticmethod
    def visualize_3d_points(
        point_data: Dict[str, List[Point3D]],
        cfg: Dict[str, Any],
        id: str,
        point_size: int = 1,
        max_points: int = 200000,
    ) -> None:
        all_points = []
        for points in point_data.values():
            all_points.extend(points)
        
        if len(all_points) == 0:
            return
        if len(all_points) > max_points:
            all_points = random.sample(all_points, max_points)

        x = [p.x for p in all_points]
        y = [p.y for p in all_points]
        z = [p.z for p in all_points]
        
        fig = go.Figure(data=[
            go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(
                    size=point_size,
                    color=z,
                    colorscale='Viridis',
                    opacity=0.8
                ),
                name='All Points'
            )
        ])
        
        x_min, x_max = min(x), max(x)
        y_min, y_max = min(y), max(y)
        z_min, z_max = -1, 8
        
        x_span = x_max - x_min
        y_span = y_max - y_min
        z_span = z_max - z_min
        
        max_span = max(x_span, y_span, z_span)

        fig.update_layout(
            scene=dict(
                xaxis_title='X (Right)',
                yaxis_title='Y (Forward)',
                zaxis_title='Z (Up)',
                aspectmode='manual',
                aspectratio=dict(
                    x=x_span/max_span,
                    y=y_span/max_span,
                    z=z_span/max_span
                ),
                camera=dict(
                    up=dict(x=0, y=0, z=1),
                    center=dict(x=0, y=0, z=0),
                    eye=dict(x=1.5, y=1.5, z=0.5)
                ),
                
                xaxis=dict(range=[x_min, x_min + max_span]),
                yaxis=dict(range=[y_min, y_min + max_span]),
                zaxis=dict(range=[z_min, z_max]),
            ),
            margin=dict(l=0, r=0, b=0, t=0)
        )
        output_folder = os.path.join(cfg.get('bev_output_folder'), cfg.get("data_name"))
        os.makedirs(output_folder, exist_ok=True)
        fig.write_html(os.path.join(output_folder, f"{id}.html"), auto_open=False)
    
    @staticmethod
    def visualize_3d_points_dynamic(
        point_data: Dict[str, List[Point3D]],
        point_size: int = 1,
        max_points: int = 10000,
    ) -> None:
        frame_dict = {}
        for key, points in point_data.items():
            try:
                frame_num = int(key.split('_')[0])
                if frame_num not in frame_dict:
                    frame_dict[frame_num] = []
                frame_dict[frame_num].extend(points)
            except (IndexError, ValueError):
                continue
        
        if not frame_dict:
            print("Warning: No valid frame data found")
            return
        
        sorted_frames = sorted(frame_dict.keys())
        
        frames = []
        total_frames = len(sorted_frames)
        
        for i, frame in enumerate(sorted_frames):
            current_points = frame_dict[frame]
            
            if len(current_points) > max_points:
                step = len(current_points) // max_points + 1
                display_points = current_points[::step][:max_points]
            else:
                display_points = current_points.copy()

            if i % 10 == 0 or i == total_frames - 1:
                print(f"Processed {i+1}/{total_frames} frames, current point count: {len(display_points)}, {display_points[0]}")
                
            x = [p.x for p in display_points]
            y = [p.y for p in display_points]
            z = [p.z for p in display_points]
            
            frames.append(
                go.Frame(
                    data=[go.Scatter3d(
                        x=x, y=y, z=z,
                        mode='markers',
                        marker=dict(
                            size=point_size,
                            color=z,
                            colorscale='Viridis',
                            opacity=0.8
                        )
                    )],
                    name=f"Frame {frame}"
                )
            )
        
        all_points_flat = [p for points in point_data.values() for p in points]
        x_all = [p.x for p in all_points_flat]
        y_all = [p.y for p in all_points_flat]
        z_all = [p.z for p in all_points_flat]
        
        max_range = max(max(x_all)-min(x_all), 
                    max(y_all)-min(y_all), 
                    max(z_all)-min(z_all))
        
        fig = go.Figure(
            data=frames[0].data,
            frames=frames,
            layout=go.Layout(
                scene=dict(
                    xaxis_title='X (Right)',
                    yaxis_title='Y (Forward)',
                    zaxis_title='Z (Up)',
                    aspectmode='manual',
                    aspectratio=dict(
                        x=1, 
                        y=(max(y_all)-min(y_all))/max_range, 
                        z=(max(z_all)-min(z_all))/max_range
                    ),
                    camera=dict(
                        up=dict(x=0, y=0, z=1),
                        center=dict(x=0, y=0, z=0),
                        eye=dict(x=1.5, y=1.5, z=0.5)
                    ),
                    xaxis=dict(range=[min(x_all), min(x_all)+max_range]),
                    yaxis=dict(range=[min(y_all), min(y_all)+max_range]), 
                    zaxis=dict(range=[min(z_all), min(z_all)+max_range])
                ),
                margin=dict(l=0, r=0, b=0, t=30),
                title = f"Point Cloud Animation ({len(sorted_frames)} frames total)",
                updatemenus=[{
                    "type": "buttons",
                    "buttons": [
                        {
                            "args": [None, {
                                "frame": {"duration": 100, "redraw": True},
                                "fromcurrent": True,
                                "transition": {"duration": 0}
                            }],
                            "label": "paly",
                            "method": "animate"
                        },
                        {
                            "args": [[None], {
                                "frame": {"duration": 0, "redraw": False},
                                "mode": "immediate"
                            }],
                            "label": "stop",
                            "method": "animate"
                        }
                    ],
                    "direction": "left",
                    "pad": {"r": 10, "t": 10},
                    "x": 0.1,
                    "xanchor": "right",
                    "y": 0,
                    "yanchor": "top"
                }],
                sliders=[{
                    "active": 0,
                    "steps": [
                        {
                            "args": [[f.name], {
                                "frame": {"duration": 0, "redraw": True},
                                "mode": "immediate"
                            }],
                            "label": f.name.split()[-1],
                            "method": "animate"
                        } for f in frames
                    ],
                    "transition": {"duration": 0},
                    "x": 0.1,
                    "len": 0.9,
                    "currentvalue": {
                        "prefix": "current: ",
                        "visible": True,
                        "xanchor": "right"
                    }
                }]
            )
        )
        
        fig.show()